package pl.edu.agh.neuraleconomy.core.nn;

import java.util.Collections;
import java.util.Comparator;
import java.util.LinkedList;
import java.util.List;

import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NonNull;
import lombok.RequiredArgsConstructor;
import lombok.Setter;

import org.apache.commons.math3.stat.descriptive.moment.Mean;
import org.apache.log4j.Logger;
import org.encog.ml.data.MLData;
import org.encog.ml.data.MLDataPair;
import org.encog.ml.data.MLDataSet;

@RequiredArgsConstructor
public class NetworkStructureOptimizer {
	private Logger logger = Logger.getLogger(getClass());

	@Setter
	@NonNull
	private Range inputLen;
	@Setter
	@NonNull
	private Range secondLayer;
	@Setter
	@NonNull
	private Range thirdLayer;
	@Setter
	@NonNull
	private Range outputLen = new Range(1, 1, 1);

	@Setter
	private double[] trainData;
	@Setter
	private double[] testData;

	public NetworkStructure getOptimalStructure(NetworkTrainer trainer) {
		List<NetworkStructure> structList = new LinkedList<NetworkStructure>();
		
		for(int inp = inputLen.getFrom(); inp <= inputLen.getTo(); inp+= inputLen.getStep()){
			for(int out = outputLen.getFrom(); out <= outputLen.getTo(); out += outputLen.getStep()){
				for(int sec = secondLayer.getFrom(); sec <= secondLayer.getTo(); sec += secondLayer.getStep()){
					for(int th = thirdLayer.getFrom(); th <= thirdLayer.getTo(); th += thirdLayer.getStep()){
						NetworkStructure structure = new NetworkStructure(inp, sec, th, out, null);
						
						SimpleNetwork network = train(trainer, inp, sec, th, out);
						Prediction prediction = test(network, structure);
						computeMeanError(structure, prediction);
						structList.add(structure);
					}
				}
			}
		}
		
		return getBestStructure(structList);
	}
	
	private NetworkStructure getBestStructure(List<NetworkStructure> structures){
		Collections.sort(structures, getStructureComparator());
		
		return structures.size() > 0 ? structures.get(0) : null;
	}
	
	private SimpleNetwork train(NetworkTrainer trainer, int input, int secondLayer, int thirdLayer, int output){
		SimpleNetwork net = new SimpleNetwork(input, secondLayer, thirdLayer, output);
		trainer.setDataSet(CoreUtils.createDataSet(trainData, input, output));
		trainer.setNetwork(net);
		trainer.train();
		return net;
	}
	
	private Prediction test(SimpleNetwork network, NetworkStructure structure){
		MLDataSet testSet = CoreUtils.createDataSet(testData, structure.getInputLen(), structure.getOutputLen());
		
		double expected [] = new double[(int) testSet.getRecordCount()];
		double predicted [] = new double[(int) testSet.getRecordCount()];
		
		for(int i=0; i < testSet.getRecordCount(); i++){
			MLDataPair pair = testSet.get(i);
			
			MLData result = network.getNetwork().compute(pair.getInput());
			
			expected[i] = pair.getIdeal().getData(0);
			predicted[i] = result.getData(0);
		}
		
		return new Prediction(predicted, expected);
	}
	
	private Comparator<NetworkStructure> getStructureComparator(){
		return new Comparator<NetworkStructure>() {

			public int compare(NetworkStructure o1, NetworkStructure o2) {
				return o1.getError().compareTo(o2.getError());
			}
		};
	}
	
	private void computeMeanError(NetworkStructure structure, Prediction prediction){
		
		double [] error = new double[prediction.predicted.length];
		
		for (int i = 0; i < error.length; i++) {
			error[i] = Math.abs(prediction.getExpected()[i] - prediction.getPredicted()[i]);
		}
		double meanError = (new Mean()).evaluate(error);
		
		logger.info(String.format("(%d, %d, %d, %d)\t%s", structure.getInputLen(), structure.getSecondLayer(), structure.getThirdLayer(), structure.getOutputLen(), String.valueOf(meanError)));
		structure.setError(meanError);	
	}

	@Data
	@AllArgsConstructor()
	public static class Range {
		int from;
		int to;
		int step;
	}

	
	@Data
	@AllArgsConstructor
	protected class Prediction {
		double [] predicted;
		double [] expected;
	}

}
